import numpy as np
from crop_coefficient_calculation import generate_Kc
from van_Genuchten_relations import volumetricMoisture

from parameters_mpc import *
# Load the  soil parameters for the MZs
soilPars_mz1 = ManagementZone_1()
soilPars_mz2 = ManagementZone_2()
soilPars_mz3 = ManagementZone_3()

# Load the weather conditions
_weather_conditions = np.loadtxt("../relevant_data/Weather Data_dry_n.txt")
_average_temps = _weather_conditions[:, 0]
_crop_coeff = generate_Kc(_average_temps)
_rain = _weather_conditions[:, 1]
_ref_Evap = _weather_conditions[:, 2]

for i in range(len(_ref_Evap)):
    if _ref_Evap[i] < 1.04:
        _ref_Evap[i] = 1.04
    if _ref_Evap[i] > 9.0:
        _ref_Evap[i] = 9.0
    pass
# Load the state trajectory over the simukation period
state_traj_mz1 = np.loadtxt('../results/state_mz1_mad40.txt')
state_traj_mz2 = np.loadtxt('../results/state_mz2_mad40.txt')
state_traj_mz3 = np.loadtxt('../results/state_mz3_mad40.txt')

# Lower and upper bounds of the target zone for each MZ
LowerZone_mz1 = 0.216
UpperZone_mz1 = 0.280
LowerZone_mz2 = 0.216
UpperZone_mz2 = 0.280
LowerZone_mz3 = 0.244
UpperZone_mz3 = 0.300

# pwp_mz1 = 0.12
# pwp_mz2 = 0.12
# pwp_mz3 = 0.16

pwp_mz1 = volumetricMoisture(-152.85, soilPars_mz1)
pwp_mz2 = volumetricMoisture(-152.85, soilPars_mz2)
pwp_mz3 = volumetricMoisture(-152.85, soilPars_mz3)

# anaerobic_point_mz1 = 0.43
# anaerobic_point_mz2 = 0.43
# anaerobic_point_mz3 = 0.39

anaerobic_point_mz1 = volumetricMoisture(-0.1, soilPars_mz1)
anaerobic_point_mz2 = volumetricMoisture(-0.1, soilPars_mz2)
anaerobic_point_mz3 = volumetricMoisture(-0.1, soilPars_mz3)

maximal_potential_yield = 8.8
yield_response_factor = 1.15


def evaluate_water_stress_factor_mz1(current_vol_moisture):

    if current_vol_moisture <= pwp_mz1:
        return 0
    elif pwp_mz1 < current_vol_moisture <= LowerZone_mz1:
        return (current_vol_moisture - pwp_mz1)/(LowerZone_mz1 - pwp_mz1)

    elif LowerZone_mz1 < current_vol_moisture < UpperZone_mz1:
        return 1

    elif UpperZone_mz1 <= current_vol_moisture <= anaerobic_point_mz1:
        return (current_vol_moisture - UpperZone_mz1)/(anaerobic_point_mz1 - UpperZone_mz1)

    else:
        return 0


def evaluate_water_stress_factor_mz2(current_vol_moisture):

    if current_vol_moisture <= pwp_mz2:
        return 0
    elif pwp_mz2 < current_vol_moisture <= LowerZone_mz2:
        return (current_vol_moisture - pwp_mz2)/(LowerZone_mz2 - pwp_mz2)

    elif LowerZone_mz2 < current_vol_moisture < UpperZone_mz2:
        return 1

    elif UpperZone_mz2 <= current_vol_moisture <= anaerobic_point_mz2:
        return (current_vol_moisture - UpperZone_mz2)/(anaerobic_point_mz2 - UpperZone_mz2)

    else:
        return 0


def evaluate_water_stress_factor_mz3(current_vol_moisture):
    if current_vol_moisture <= pwp_mz3:
        return 0
    elif pwp_mz3 < current_vol_moisture <= LowerZone_mz3:
        return (current_vol_moisture - pwp_mz3)/(LowerZone_mz3 - pwp_mz3)

    elif LowerZone_mz3 < current_vol_moisture < UpperZone_mz3:
        return 1

    elif UpperZone_mz3 <= current_vol_moisture <= anaerobic_point_mz3:
        return (current_vol_moisture - UpperZone_mz3)/(anaerobic_point_mz3 - UpperZone_mz3)

    else:
        return 0


water_stress_factor_mz1 = []
water_stress_factor_mz2 = []
water_stress_factor_mz3 = []

maximum_ET_mz1 = []
maximum_ET_mz2 = []
maximum_ET_mz3 = []

seasonal_ET_mz1 = []
seasonal_ET_mz2 = []
seasonal_ET_mz3 = []


for i in range(len(state_traj_mz1)):
    water_stress_factor_mz1.append(
        evaluate_water_stress_factor_mz1(state_traj_mz1[i]))
    water_stress_factor_mz2.append(
        evaluate_water_stress_factor_mz2(state_traj_mz2[i]))
    water_stress_factor_mz3.append(
        evaluate_water_stress_factor_mz3(state_traj_mz3[i]))

    maximum_ET_mz1.append(_crop_coeff[i]*_ref_Evap[i])
    maximum_ET_mz2.append(_crop_coeff[i]*_ref_Evap[i])
    maximum_ET_mz3.append(_crop_coeff[i]*_ref_Evap[i])

    seasonal_ET_mz1.append(evaluate_water_stress_factor_mz1(
        state_traj_mz1[i])*_crop_coeff[i]*_ref_Evap[i])
    seasonal_ET_mz2.append(evaluate_water_stress_factor_mz2(
        state_traj_mz2[i])*_crop_coeff[i]*_ref_Evap[i])
    seasonal_ET_mz3.append(evaluate_water_stress_factor_mz3(
        state_traj_mz3[i])*_crop_coeff[i]*_ref_Evap[i])
    pass

water_stress_factor_mz1 = np.array(water_stress_factor_mz1)
water_stress_factor_mz2 = np.array(water_stress_factor_mz2)
water_stress_factor_mz3 = np.array(water_stress_factor_mz3)


# yield_reduction_factor_mz1 = yield_response_factor *sum(1-water_stress_factor_mz1)
# yield_reduction_factor_mz2 = yield_response_factor *sum(1-water_stress_factor_mz2)
# yield_reduction_factor_mz3 = yield_response_factor *sum(1-water_stress_factor_mz3)
# predicted_yield_mz1 = maximal_potential_yield*(1-yield_reduction_factor_mz1)
# predicted_yield_mz2 = maximal_potential_yield*(1-yield_reduction_factor_mz2)
# predicted_yield_mz3 = maximal_potential_yield*(1-yield_reduction_factor_mz3)

pred_yield_mz1 = (1 - yield_response_factor + yield_response_factor *
                  (sum(np.array(seasonal_ET_mz1))/sum(np.array(maximum_ET_mz1))))*maximal_potential_yield
pred_yield_mz2 = (1 - yield_response_factor + yield_response_factor *
                  (sum(np.array(seasonal_ET_mz2))/sum(np.array(maximum_ET_mz2))))*maximal_potential_yield
pred_yield_mz3 = (1 - yield_response_factor + yield_response_factor *
                  (sum(np.array(seasonal_ET_mz3))/sum(np.array(maximum_ET_mz3))))*maximal_potential_yield

average_yield = (pred_yield_mz1 + pred_yield_mz2 + pred_yield_mz3)/3
